#include <algorithm>
#include <cassert>
#include "BVH.hpp"

BVHAccel::BVHAccel(std::vector<Object*> p, int maxPrimsInNode,
                   SplitMethod splitMethod)
    : maxPrimsInNode(std::min(255, maxPrimsInNode)), splitMethod(splitMethod),
      primitives(std::move(p))
{
    time_t start, stop;
    time(&start);
    if (primitives.empty())
        return;

    root = recursiveBuild(primitives);

    time(&stop);
    double diff = difftime(stop, start);
    int hrs = (int)diff / 3600;
    int mins = ((int)diff / 60) - (hrs * 60);
    int secs = (int)diff - (hrs * 3600) - (mins * 60);

    printf(
        "\rBVH Generation complete: \nTime Taken: %i hrs, %i mins, %i secs\n\n",
        hrs, mins, secs);
}

// SAH
BVHBuildNode* BVHAccel::recursiveBuild(std::vector<Object*> objects)
{
    BVHBuildNode* node = new BVHBuildNode();

    // Compute bounds of all primitives in BVH node
    Bounds3 bounds;
    for (int i = 0; i < objects.size(); ++i)
        bounds = Union(bounds, objects[i]->getBounds());
    if (objects.size() == 1) {
        // Create leaf _BVHBuildNode_
        node->bounds = objects[0]->getBounds();
        node->object = objects[0];
        node->left = nullptr;
        node->right = nullptr;
        return node;
    }
    else if (objects.size() == 2) {
        node->left = recursiveBuild(std::vector{objects[0]});
        node->right = recursiveBuild(std::vector{objects[1]});

        node->bounds = Union(node->left->bounds, node->right->bounds);
        return node;
    }
    else {
        Bounds3 centroidBounds;
        for (int i = 0; i < objects.size(); ++i)
            centroidBounds = Union(centroidBounds, objects[i]->getBounds().Centroid());
        // NAIVE
        if (splitMethod == SplitMethod::NAIVE) {
            int dim = centroidBounds.maxExtent();
            switch (dim) {
                case 0:
                    std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                        return f1->getBounds().Centroid().x <
                            f2->getBounds().Centroid().x;
                    });
                    break;
                case 1:
                    std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                        return f1->getBounds().Centroid().y <
                            f2->getBounds().Centroid().y;
                    });
                    break;
                case 2:
                    std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                        return f1->getBounds().Centroid().z <
                            f2->getBounds().Centroid().z;
                    });
                    break;
            }

            auto beginning = objects.begin();
            auto middling = objects.begin() + (objects.size() / 2);
            auto ending = objects.end();

            auto leftshapes = std::vector<Object*>(beginning, middling);
            auto rightshapes = std::vector<Object*>(middling, ending);

            assert(objects.size() == (leftshapes.size() + rightshapes.size()));

            node->left = recursiveBuild(leftshapes);
            node->right = recursiveBuild(rightshapes);

            node->bounds = Union(node->left->bounds, node->right->bounds);
        }
        // SAH
        else if (splitMethod == SplitMethod::SAH) {
            //SAH相关参数：
            float Sn = centroidBounds.SurfaceArea(); //Sn
            int B = 10; //B usually < 32;
            int minCostCoor = 0, mincostIndex = 0; //最优分割方案，轴和B的值
            float minCost = std::numeric_limits<float>::infinity(); //最小花费

            //分别进行三个轴的划分
            for (int i = 0; i < 3; i++) {   //对三角形按照所选择的轴的位置进行排序
                switch (i) {
                    case 0:  //x轴最长，对x进行排序
                        std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {return f1->getBounds().Centroid().x < f2->getBounds().Centroid().x;});
                        break;
                    case 1:  //y轴最长，对y进行排序
                        std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {return f1->getBounds().Centroid().y < f2->getBounds().Centroid().y;});
                        break;
                    case 2:  //z轴最长，对z进行排序
                        std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {return f1->getBounds().Centroid().z < f2->getBounds().Centroid().z;});
                        break;
                }
                for(int j = 1; j < B; j++) {
                    auto beginning = objects.begin();
                    auto middling = objects.begin() + (objects.size() * j / B);
                    auto ending = objects.end();
                    auto leftshapes = std::vector<Object*>(beginning, middling);
                    auto rightshapes = std::vector<Object*>(middling, ending);
                    //求左右包围盒:
                    Bounds3 leftBounds, rightBounds;
                    for (int k = 0; k < leftshapes.size(); ++k)
                        leftBounds = Union(leftBounds, leftshapes[k]->getBounds().Centroid());
                    for (int k = 0; k < rightshapes.size(); ++k)
                        rightBounds = Union(rightBounds, rightshapes[k]->getBounds().Centroid());
                    float SA = leftBounds.SurfaceArea(); //SA
                    float SB = rightBounds.SurfaceArea(); //SB
                    float cost = 1 + (leftshapes.size() * SA + rightshapes.size() * SB) / Sn; //计算花费
                    if(cost < minCost) //如果花费更小，记录当前坐标值
                    {
                        minCost = cost;
                        mincostIndex = j;
                        minCostCoor = i;
                    }
                }
            }
            switch (minCostCoor) {
                case 0:  //对x进行排序
                    std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {return f1->getBounds().Centroid().x < f2->getBounds().Centroid().x;});
                    break;
                case 1:  //对y进行排序
                    std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {return f1->getBounds().Centroid().y < f2->getBounds().Centroid().y;});
                    break;
                case 2:  //对z进行排序
                    std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {return f1->getBounds().Centroid().z < f2->getBounds().Centroid().z;});
                    break;
            }
            auto beginning = objects.begin();
            auto middling = objects.begin() + (objects.size() * mincostIndex / B);
            auto ending = objects.end();
            auto leftshapes = std::vector<Object*>(beginning, middling);
            auto rightshapes = std::vector<Object*>(middling, ending);
            assert(objects.size() == (leftshapes.size() + rightshapes.size()));
            node->left = recursiveBuild(leftshapes);
            node->right = recursiveBuild(rightshapes);
            node->bounds = Union(node->left->bounds, node->right->bounds);
        }
    }
    return node;
}

Intersection BVHAccel::Intersect(const Ray& ray) const
{
    Intersection isect;
    if (!root)
        return isect;
    isect = BVHAccel::getIntersection(root, ray);
    return isect;
}

Intersection BVHAccel::getIntersection(BVHBuildNode* node, const Ray& ray) const
{
    // TODO Traverse the BVH to find intersection
    Intersection inter;
    Vector3f invdir(1 / ray.direction.x, 1 / ray.direction.y, 1 / ray.direction.z);
    std::array<int, 3> dirIsNeg;
    dirIsNeg[0] = ray.direction.x < 0;
    dirIsNeg[1] = ray.direction.y < 0;
    dirIsNeg[2] = ray.direction.z < 0;

    if (!node->bounds.IntersectP(ray, invdir, dirIsNeg)) return inter;
    if(node -> left == nullptr && node -> right == nullptr) return node->object->getIntersection(ray);

    Intersection hit1 = getIntersection(node -> left,  ray);
    Intersection hit2 = getIntersection(node -> right, ray);

    return hit1.distance < hit2.distance ? hit1 : hit2;
}